import os
os.environ['ETS_TOOLKIT'] = 'qt4'
from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine, PipelineBase
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
import numpy as np
from numpy import array, random, linspace, pi, ravel, cos, sin, empty, mgrid, arange
from mayavi import mlab
import networkx as nx
from tvtk.api import tvtk
from mayavi.sources.builtin_surface import BuiltinSurface
import csv
from scipy.special import sph_harm

engine = Engine()
engine.start()


class Form(HasTraits):
    atomic_orbital = Instance(MlabSceneModel, args={'engine': engine})

    @observe('atomic_orbital.activated')
    def update_atomic_orbital(self, event=None):
        mlab = self.atomic_orbital.mlab
        x, y, z = np.ogrid[- .5:.5:200j, - .5:.5:200j, - .5:.5:200j]
        r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
        # Generalized Laguerre polynomial (3, 2)
        L = - r ** 3 / 6 + 5. / 2 * r ** 2 - 10 * r + 6
        # Spherical harmonic (3, 2)
        Y = (x + y * 1j) ** 2 * z / r ** 3
        Phi = L * Y * np.exp(- r) * r ** 2
        # Plot it ####################################################################
        # We create a scalar field with the module of Phi as the scalar
        src = mlab.pipeline.scalar_field(np.abs(Phi))
        src.image_data.point_data.add_array(np.angle(Phi).T.ravel())
        # We need to give a name to our new dataset.
        src.image_data.point_data.get_array(1).name = 'angle'
        # Make sure that the dataset is up to date with the different arrays:
        src.update()
        # We select the 'scalar' attribute, ie the norm of Phi
        src2 = mlab.pipeline.set_active_attribute(src, point_scalars='scalar')
        # Cut isosurfaces of the norm
        contour = mlab.pipeline.contour(src2)
        # Now we select the 'angle' attribute, ie the phase of Phi
        contour2 = mlab.pipeline.set_active_attribute(contour, point_scalars='angle')
        # And we display the surface. The colormap is the current attribute: the phase.
        mlab.pipeline.surface(contour2, colormap='hsv')
        # self.scene.mlab.colorbar(title='Phase', orientation='vertical', nb_labels=3)

    boy = Instance(MlabSceneModel, args={'engine': engine})

    @observe('boy.activated')
    def update_boy(self, event=None):
        mlab = self.boy.mlab
        u, v = np.mgrid[- 0.035:np.pi:0.01, - 0.05:np.pi:0.01]
        X = 2 / 3. * (np.cos(u) * np.cos(2 * v) + np.sqrt(2) * np.sin(u) * np.cos(v)) * np.cos(u) / (np.sqrt(2) - np.sin(2 * u) * np.sin(3 * v))
        Y = 2 / 3. * (np.cos(u) * np.sin(2 * v) - np.sqrt(2) * np.sin(u) * np.sin(v)) * np.cos(u) / (np.sqrt(2) - np.sin(2 * u) * np.sin(3 * v))
        Z = -np.sqrt(2) * np.cos(u) * np.cos(u) / (np.sqrt(2) - np.sin(2 * u) * np.sin(3 * v))
        S = np.sin(u)
        mlab.mesh(X, Y, Z, scalars=S, colormap='YlGnBu', )

    polydata = Instance(MlabSceneModel, args={'engine': engine})

    @observe('polydata.activated')
    def update_polydata(self, event=None):
        mlab = self.polydata.mlab

        def polydata():
            # The numpy array data.
            points = np.array([[0, -0.5, 0], [1.5, 0, 0], [0, 1, 0], [0, 0, 0.5], [-1, -1.5, 0.1], [0, -1, 0.5], [-1, -0.5, 0], [1, 0.8, 0]], 'f')
            triangles = np.array([[0, 1, 3], [1, 2, 3], [1, 0, 5], [2, 3, 4], [3, 0, 4], [0, 5, 4], [2, 4, 6], [2, 1, 7]])
            scalars = np.random.random(points.shape)
            # The TVTK dataset.
            mesh = tvtk.PolyData(points=points, polys=triangles)
            mesh.point_data.scalars = scalars
            mesh.point_data.scalars.name = 'scalars'
            return mesh

        def view(dataset):
            surf = mlab.pipeline.surface(dataset, opacity=0.1)
            mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

        view(polydata())

    polydata1 = Instance(MlabSceneModel, args={'engine': engine})

    @observe('polydata1.activated')
    def update_polydata1(self, event=None):
        mlab = self.polydata1.mlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        # scalars = np.array([[0, 1, 0]])
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.point_data.scalars = scalars
        # mesh.cell_data.scalars = scalars
        surf = mlab.pipeline.surface(mesh)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    polydata3 = Instance(MlabSceneModel, args={'engine': engine})

    @observe('polydata3.activated')
    def update_polydata3(self, event=None):
        mlab = self.polydata3.mlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([[0, .5, 1]])
        mesh = tvtk.PolyData(points=points, polys=triangles)
        mesh.cell_data.scalars = scalars
        surf = mlab.pipeline.surface(mesh)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )


    polydata2 = Instance(MlabSceneModel, args={'engine': engine})

    @observe('polydata2.activated')
    def update_polydata2(self, event=None):
        mlab = self.polydata2.mlab
        points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [2, 0, 0]], 'f')
        triangles = np.array([[0, 1, 2], [1, 3, 2], [1, 4, 3], ])
        scalars = np.array([0, .1, .2, .3, .4], dtype=float)
        surf = mlab.triangular_mesh(points[:, 0], points[:, 1], points[:, 2], triangles, scalars=scalars)
        mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

    image_data = Instance(MlabSceneModel, args={'engine': engine})

    @observe('image_data.activated')
    def update_image_data(self, event=None):
        mlab = self.image_data.mlab

        def image_data():
            data = np.random.random((3, 3, 3))
            i = tvtk.ImageData(spacing=(1, 1, 1), origin=(0, 0, 0))
            i.point_data.scalars = data.ravel()
            i.point_data.scalars.name = 'scalars'
            i.dimensions = data.shape
            return i

        def view(dataset):
            surf = mlab.pipeline.surface(dataset, opacity=0.1)
            mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

        view(image_data())

    rectilinear_grid = Instance(MlabSceneModel, args={'engine': engine})

    @observe('rectilinear_grid.activated')
    def update_rectilinear_grid(self, event=None):
        mlab = self.rectilinear_grid.mlab

        def rectilinear_grid():
            data = np.random.random((3, 3, 3))
            r = tvtk.RectilinearGrid()
            r.point_data.scalars = data.ravel()
            r.point_data.scalars.name = 'scalars'
            r.dimensions = data.shape
            r.x_coordinates = np.array((0, 0.7, 1.4))
            r.y_coordinates = np.array((0, 1, 3))
            r.z_coordinates = np.array((0, .5, 2))
            return r

        def view(dataset):
            surf = mlab.pipeline.surface(dataset, opacity=0.1)
            mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

        view(rectilinear_grid())

    structured_grid = Instance(MlabSceneModel, args={'engine': engine})

    @observe('structured_grid.activated')
    def update_structured_grid(self, event=None):
        mlab = self.structured_grid.mlab

        def generate_annulus(r, theta, z):
            """ Generate points for structured grid for a cylindrical annular
                volume.  This method is useful for generating a unstructured
                cylindrical mesh for VTK (and perhaps other tools).
            """
            # Find the x values and y values for each plane.
            x_plane = (np.cos(theta) * r[:, None]).ravel()
            y_plane = (np.sin(theta) * r[:, None]).ravel()
            # Allocate an array for all the points.  We'll have len(x_plane)
            # points on each plane, and we have a plane for each z value, so
            # we need len(x_plane)*len(z) points.
            points = np.empty([len(x_plane) * len(z), 3])
            # Loop through the points for each plane and fill them with the
            # correct x,y,z values.
            start = 0
            for z_plane in z:
                end = start + len(x_plane)
                # slice out a plane of the output points and fill it
                # with the x,y, and z values for this plane.  The x,y
                # values are the same for every plane.  The z value
                # is set to the current z
                plane_points = points[start:end]
                plane_points[:, 0] = x_plane
                plane_points[:, 1] = y_plane
                plane_points[:, 2] = z_plane
                start = end
            return points

        def structured_grid():
            # Make the data.
            dims = (3, 4, 3)
            r = np.linspace(5, 15, dims[0])
            theta = np.linspace(0, 0.5 * np.pi, dims[1])
            z = np.linspace(0, 10, dims[2])
            pts = generate_annulus(r, theta, z)
            sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
            sgrid.points = pts
            s = np.random.random((dims[0] * dims[1] * dims[2]))
            sgrid.point_data.scalars = np.ravel(s.copy())
            sgrid.point_data.scalars.name = 'scalars'
            return sgrid

        def view(dataset):
            surf = mlab.pipeline.surface(dataset, opacity=0.1)
            mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

        view(structured_grid())

    unstructured_grid = Instance(MlabSceneModel, args={'engine': engine})

    @observe('unstructured_grid.activated')
    def update_unstructured_grid(self, event=None):
        mlab = self.unstructured_grid.mlab

        def generate_annulus(r, theta, z):
            """ Generate points for structured grid for a cylindrical annular
                volume.  This method is useful for generating a unstructured
                cylindrical mesh for VTK (and perhaps other tools).
            """
            # Find the x values and y values for each plane.
            x_plane = (np.cos(theta) * r[:, None]).ravel()
            y_plane = (np.sin(theta) * r[:, None]).ravel()
            # Allocate an array for all the points.  We'll have len(x_plane)
            # points on each plane, and we have a plane for each z value, so
            # we need len(x_plane)*len(z) points.
            points = np.empty([len(x_plane) * len(z), 3])
            # Loop through the points for each plane and fill them with the
            # correct x,y,z values.
            start = 0
            for z_plane in z:
                end = start + len(x_plane)
                # slice out a plane of the output points and fill it
                # with the x,y, and z values for this plane.  The x,y
                # values are the same for every plane.  The z value
                # is set to the current z
                plane_points = points[start:end]
                plane_points[:, 0] = x_plane
                plane_points[:, 1] = y_plane
                plane_points[:, 2] = z_plane
                start = end
            return points

        def unstructured_grid():
            points = np.array([[0, 1.2, 0.6], [1, 0, 0], [0, 1, 0], [1, 1, 1],  # tetra
                               [1, 0, -0.5], [2, 0, 0], [2, 1.5, 0], [0, 1, 0],
                               [1, 0, 0], [1.5, -0.2, 1], [1.6, 1, 1.5], [1, 1, 1],  # Hex
                               ], 'f')
            # The cells
            cells = np.array([4, 0, 1, 2, 3,  # tetra
                              8, 4, 5, 6, 7, 8, 9, 10, 11  # hex
                              ])
            # The offsets for the cells, i.e. the indices where the cells
            # start.
            offset = np.array([0, 5])
            tetra_type = tvtk.Tetra().cell_type  # VTK_TETRA == 10
            hex_type = tvtk.Hexahedron().cell_type  # VTK_HEXAHEDRON == 12
            cell_types = np.array([tetra_type, hex_type])
            # Create the array of cells unambiguously.
            cell_array = tvtk.CellArray()
            cell_array.set_cells(2, cells)
            # Now create the UG.
            ug = tvtk.UnstructuredGrid(points=points)
            # Now just set the cell types and reuse the ug locations and cells.
            ug.set_cells(cell_types, offset, cell_array)
            scalars = np.random.random(points.shape[0])
            ug.point_data.scalars = scalars
            ug.point_data.scalars.name = 'scalars'
            return ug

        def view(dataset):
            surf = mlab.pipeline.surface(dataset, opacity=0.1)
            mlab.pipeline.surface(mlab.pipeline.extract_edges(surf), color=(0, 0, 0), )

        view(unstructured_grid())

    delaunay_graph = Instance(MlabSceneModel, args={'engine': engine})

    @observe('delaunay_graph.activated')
    def update_delaunay_graph(self, event=None):
        mlab = self.delaunay_graph.mlab

        def compute_delaunay_edges(x, y, z, visualize=False):
            """ Given 3-D points, returns the edges of their
                Delaunay triangulation.
                Parameters
                -----------
                x: ndarray
                    x coordinates of the points
                y: ndarray
                    y coordinates of the points
                z: ndarray
                    z coordinates of the points
                Returns
                ---------
                new_x: ndarray
                    new x coordinates of the points (same coords but different
                    assignment of points)
                new_y: ndarray
                    new y coordinates of the points (same coords but different
                    assignment of points)
                new_z: ndarray
                    new z coordinates of the points (same coords but different
                    assignment of points)
                edges: 2D ndarray.
                    The indices of the edges of the Delaunay triangulation as a
                    (N, 2) array [[pair1_index1, pair1_index2],
                                  [pair2_index1, pair2_index2],
                                  ...                         ]
            """
            if visualize:
                vtk_source = mlab.points3d(x, y, z, opacity=0.3, mode='2dvertex')
                vtk_source.actor.property.point_size = 3
            else:
                vtk_source = mlab.pipeline.scalar_scatter(x, y, z, figure=False)
            delaunay = mlab.pipeline.delaunay3d(vtk_source)
            delaunay.filter.offset = 999  # seems more reliable than the default
            edges = mlab.pipeline.extract_edges(delaunay)
            if visualize:
                mlab.pipeline.surface(edges, opacity=0.3, line_width=3)
            # We extract the output array. the 'points' attribute itself
            # is a TVTK array, that we convert to a numpy array using
            # its 'to_array' method.
            dataset = edges.get_output_dataset()
            new_x, new_y, new_z = dataset.points.to_array().T
            lines = dataset.lines.to_array()
            return new_x, new_y, new_z, np.array([lines[1::3], lines[2::3]]).T

        def graph_plot(x, y, z, start_idx, end_idx, edge_scalars=None, **kwargs):
            """ Show the graph edges using Mayavi
                Parameters
                -----------
                x: ndarray
                    x coordinates of the points
                y: ndarray
                    y coordinates of the points
                z: ndarray
                    z coordinates of the points
                edge_scalars: ndarray, optional
                    optional data to give the color of the edges.
                kwargs:
                    extra keyword arguments are passed to quiver3d.
            """
            vec = mlab.quiver3d(x[start_idx], y[start_idx], z[start_idx],
                                x[end_idx] - x[start_idx], y[end_idx] - y[start_idx], z[end_idx] - z[start_idx],
                                scalars=edge_scalars, mode='2ddash', scale_factor=1, **kwargs)
            if edge_scalars is not None:
                vec.glyph.color_mode = 'color_by_scalar'
            return vec

        def build_geometric_graph(x, y, z, edges):
            """ Build a NetworkX graph with xyz node coordinates and the node indices
                of the end nodes.
                Parameters
                -----------
                x: ndarray
                    x coordinates of the points
                y: ndarray
                    y coordinates of the points
                z: ndarray
                    z coordinates of the points
                edges: the (2, N) array returned by compute_delaunay_edges()
                    containing node indices of the end nodes. Weights are applied to
                    the edges based on their euclidean length for use by the MST
                    algorithm.
                Returns
                ---------
                g: A NetworkX undirected graph
                Notes
                ------
                We don't bother putting the coordinates into the NX graph.
                Instead the graph node is an index to the column.
            """
            xyz = np.array((x, y, z))

            def euclidean_dist(i, j):
                d = xyz[:, i] - xyz[:, j]
                return np.sqrt(np.dot(d, d))

            g = nx.Graph()
            for i, j in edges:
                if nx.__version__.split('.')[0] > '0':
                    g.add_edge(i, j, weight=euclidean_dist(i, j))
                else:
                    g.add_edge(i, j, euclidean_dist(i, j))
            return g

        def points_on_sphere(N):
            """ Generate N evenly distributed points on the unit sphere centered at
                the origin. Uses the 'Golden Spiral'.
                Code by Chris Colbert from the numpy-discussion list.
            """
            phi = (1 + np.sqrt(5)) / 2  # the golden ratio
            long_incr = 2 * np.pi / phi  # how much to increment the longitude
            dz = 2.0 / float(N)  # a unit sphere has diameter 2
            bands = np.arange(N)  # each band will have one point placed on it
            z = bands * dz - 1 + (dz / 2)  # the height z of each band/point
            r = np.sqrt(1 - z * z)  # project onto xy-plane
            az = bands * long_incr  # azimuthal angle of point modulo 2 pi
            x = r * np.cos(az)
            y = r * np.sin(az)
            return x, y, z

        # ################################################################################
        # if __name__ == '__main__':
        # generate some points
        x, y, z = points_on_sphere(50)
        # Avoid triangulation problems on the sphere
        z *= 1.01
        # mlab.figure(1, bgcolor=(0, 0, 0))
        # mlab.clf()
        # Now get the Delaunay Triangulation from vtk via mayavi mlab. Vtk stores
        # its points in a different order so overwrite ours to match the edges
        new_x, new_y, new_z, edges = compute_delaunay_edges(x, y, z, visualize=True)
        assert (x.shape == new_x.shape)  # check triangulation got everything
        x, y, z = new_x, new_y, new_z
        if nx.__version__ < '0.99':
            raise ImportError('The version of NetworkX must be at least '
                              '0.99 to run this example')
        # Make a NetworkX graph out of our point and edge data
        g = build_geometric_graph(x, y, z, edges)
        # Compute minimum spanning tree using networkx
        # nx.mst returns an edge generator
        edges = nx.minimum_spanning_tree(g).edges(data=True)
        start_idx, end_idx, _ = np.array(list(edges)).T
        start_idx = start_idx.astype(int)
        end_idx = end_idx.astype(int)
        # Plot this with Mayavi
        graph_plot(x, y, z, start_idx, end_idx, edge_scalars=z[start_idx], opacity=0.8, colormap='summer', line_width=4, )

    custom_colormap = Instance(MlabSceneModel, args={'engine': engine})

    @observe('custom_colormap.activated')
    def update_custom_colormap(self, event=None):
        mlab = self.custom_colormap.mlab
        x, y = np.mgrid[-10:10:200j, -10:10:200j]
        z = 100 * np.sin(x * y) / (x * y)
        # self.scene.mlab.figure(bgcolor=(1, 1, 1))
        surf = mlab.surf(z, colormap='cool')
        lut = surf.module_manager.scalar_lut_manager.lut.table.to_array()
        lut[:, -1] = np.linspace(0, 255, 256)
        surf.module_manager.scalar_lut_manager.lut.table = lut

    fluidArrow_grid = Instance(MlabSceneModel, args={'engine': engine})

    @observe('fluidArrow_grid.activated')
    def update_fluidArrow_grid(self, event=None):
        x, y, z = np.mgrid[1:6:11j, 0:4:13j, 0:3:6j]
        base = x[..., 0] + y[..., 0]
        for i in range(z.shape[2]):
            z[..., i] = base * 0.25 * i
        pts = np.empty(z.shape + (3,), dtype=float)
        pts[..., 0] = x
        pts[..., 1] = y
        pts[..., 2] = z
        scalars = x * x + y * y + z * z
        vectors = np.empty(z.shape + (3,), dtype=float)
        vectors[..., 0] = (4 - y * 2)
        vectors[..., 1] = (x * 3 - 12)
        vectors[..., 2] = np.sin(z * np.pi)
        pts = pts.transpose(2, 1, 0, 3).copy()
        pts.shape = pts.size // 3, 3
        scalars = scalars.T.copy()
        vectors = vectors.transpose(2, 1, 0, 3).copy()
        vectors.shape = vectors.size // 3, 3
        sg = tvtk.StructuredGrid(dimensions=x.shape, points=pts)
        sg.point_data.scalars = scalars.ravel()
        sg.point_data.scalars.name = 'temperature'
        sg.point_data.vectors = vectors
        sg.point_data.vectors.name = 'velocity'
        d = mlab.pipeline.add_dataset(sg)
        gx = mlab.pipeline.grid_plane(d)
        gy = mlab.pipeline.grid_plane(d)
        gy.grid_plane.axis = 'y'
        gz = mlab.pipeline.grid_plane(d)
        gz.grid_plane.axis = 'z'
        iso = mlab.pipeline.iso_surface(d)
        iso.contour.maximum_contour = 75.0
        vec = mlab.pipeline.vectors(d)
        vec.glyph.mask_input_points = True
        vec.glyph.glyph.scale_factor = 1.5

    image_cursor_filter = Instance(MlabSceneModel, args={'engine': engine})

    @observe('image_cursor_filter.activated')
    def update_image_cursor_filter(self, event=None):
        x, y, z = np.ogrid[0:1:40j, 0:1:40j, 0:1:40j]
        sig = 0.5
        center = 0.5
        g = np.exp(-((x - center) ** 2 + (y - center) ** 2 + (z - center) ** 2) / (2 * sig ** 2))
        s = mlab.pipeline.scalar_field(g)
        cursor = mlab.pipeline.user_defined(s, filter='ImageCursor3D')
        cursor.filter.cursor_position = np.array([20, 20, 20])
        cursor.filter.cursor_value = 0
        cursor.filter.cursor_radius = 10
        ipw = mlab.pipeline.image_plane_widget(cursor, plane_orientation='x_axes', slice_index=20)
        mlab.colorbar()
        mlab.view(15, 70, 100, [20, 20, 20])

    julia_set = Instance(MlabSceneModel, args={'engine': engine})

    @observe('julia_set.activated')
    def update_julia_set(self, event=None):
        mlab = self.julia_set.mlab
        x, y = np.ogrid[-1.5:0.5:500j, -1:1:500j]
        z = x + 1j * y
        julia = np.zeros(z.shape)
        for i in range(50):
            z = z ** 2 - 0.70176 - 0.3842j
            julia += 1 / float(2 + i) * (z * np.conj(z) > 4)
        mlab.surf(julia, colormap='gist_earth', warp_scale='auto', vmax=1.5)

    lorenz = Instance(MlabSceneModel, args={'engine': engine})

    @observe('lorenz.activated')
    def update_lorenz(self, event=None):
        mlab = self.lorenz.mlab

        def lorenz(x, y, z, s=10., r=28., b=8. / 3.):
            """The Lorenz system."""
            u = s * (y - x)
            v = r * x - y - x * z
            w = x * y - b * z
            return u, v, w

        x, y, z = np.mgrid[-50:50:100j, -50:50:100j, -10:60:70j]
        u, v, w = lorenz(x, y, z)
        # fig = mlab.figure(size=(400, 300), bgcolor=(0, 0, 0))
        f = mlab.flow(x, y, z, u, v, w, line_width=3, colormap='Paired')
        f.module_manager.scalar_lut_manager.reverse_lut = True
        f.stream_tracer.integration_direction = 'both'
        # f.stream_tracer.maximum_propagation = 200
        src = f.mlab_source.m_data
        e = mlab.pipeline.extract_vector_components(src)
        e.component = 'z-component'
        zc = mlab.pipeline.iso_surface(e, opacity=0.5, contours=[0, ], color=(0.6, 1, 0.2))
        zc.actor.property.backface_culling = True

    RCS_3D = Instance(MlabSceneModel, args={'engine': engine})

    @observe('RCS_3D.activated')
    def update_RCS_3D(self, event=None):
        mlab = self.RCS_3D.mlab
        N = 300
        t = np.linspace(-2 * np.pi, 2 * np.pi, N)
        x = list()
        y = list()
        z = list()
        s = list()
        connections = list()
        index = 0
        for i in range(50):
            x.append(np.sin(t))
            y.append(np.cos((2 + .02 * i) * t))
            z.append(np.cos((3 + .02 * i) * t))
            s.append(t)
            connections.append(np.vstack(
                [np.arange(index, index + N - 1.5), np.arange(index + 1, index + N - .5)]).T)
            index += N
        x = np.hstack(x)
        y = np.hstack(y)
        z = np.hstack(z)
        s = np.hstack(s)
        connections = np.vstack(connections)
        src = mlab.pipeline.scalar_scatter(x, y, z, s)
        src.mlab_source.dataset.lines = connections
        src.update()
        lines = mlab.pipeline.stripper(src)
        mlab.pipeline.surface(lines, colormap='Accent', line_width=1, opacity=.4)

    select_red_balls = Instance(MlabSceneModel, args={'engine': engine})

    @observe('select_red_balls.activated')
    def update_select_red_balls(self, event=None):
        mlab = self.select_red_balls.mlab
        figure = mlab.gcf()
        # mlab.clf()
        # figure.scene.disable_render = True
        x1, y1, z1 = np.random.random((3, 10))
        red_glyphs = mlab.points3d(x1, y1, z1, color=(1, 0, 0), resolution=20)
        x2, y2, z2 = np.random.random((3, 10))
        white_glyphs = mlab.points3d(x2, y2, z2, color=(0.9, 0.9, 0.9), resolution=20)
        outline = mlab.outline(line_width=3)
        outline.outline_mode = 'cornered'
        outline.bounds = (x1[0] - 0.1, x1[0] + 0.1, y1[0] - 0.1, y1[0] + 0.1, z1[0] - 0.1, z1[0] + 0.1)
        # figure.scene.disable_render = False
        glyph_points = red_glyphs.glyph.glyph_source.glyph_source.output.points.to_array()

        def picker_callback(picker):
            if picker.actor in red_glyphs.actor.actors:
                point_id = picker.point_id // glyph_points.shape[0]
                if point_id != -1:
                    x, y, z = x1[point_id], y1[point_id], z1[point_id]
                    outline.bounds = (x - 0.1, x + 0.1, y - 0.1, y + 0.1, z - 0.1, z + 0.1)

        picker = figure.on_mouse_pick(picker_callback)
        picker.tolerance = 0.01
        mlab.title('Click on red balls')

    spherical_harmonics = Instance(MlabSceneModel, args={'engine': engine})

    @observe('spherical_harmonics.activated')
    def update_spherical_harmonics(self, event=None):
        mlab = self.spherical_harmonics.mlab
        r = 0.3
        phi, theta = np.mgrid[0:np.pi:101j, 0:2 * np.pi:101j]
        x = r * np.sin(phi) * np.cos(theta)
        y = r * np.sin(phi) * np.sin(theta)
        z = r * np.cos(phi)
        # Represent spherical harmonics on the surface of the sphere
        for n in range(0, 5):
            for m in range(n):
                s = sph_harm(m, n, theta, phi).real
                mlab.mesh(x + m, y + n, z, scalars=s, colormap='jet')
                s[s < 0] *= 0.97
                s /= s.max()
                mlab.mesh(s * x + m, s * y + n, s * z + 1.3, scalars=s, colormap='jet')

    three3d_wigner = Instance(MlabSceneModel, args={'engine': engine})

    @observe('three3d_wigner.activated')
    def update_three3d_wigner(self, event=None):
        mlab = self.three3d_wigner.mlab

        def cat(x, y, alpha, eta=1, purity=1):
            return (1 + eta * (np.exp(-x ** 2 - (y - alpha) ** 2) + np.exp(-x ** 2 - (y + alpha) ** 2) + 2 * purity * np.exp(
                -x ** 2 - y ** 2) * np.cos(2 * alpha * x)) / (2 * (1 + np.exp(- alpha ** 2)))) / 2

        x, y = np.mgrid[-4:4.15:0.1, -4:4.15:0.1]
        cat1 = cat(x, y, 1)
        cat2 = cat(x, y, 2)
        cat3 = cat(x, y, 3)
        cat1 -= 0.5
        cat2 -= 0.5
        cat3 -= 0.5
        cat1_extent = (-14, -6, -4, 4, 0, 5)
        surf_cat1 = mlab.surf(x - 10, y, cat1, colormap='Spectral', warp_scale=5, extent=cat1_extent, vmin=-0.5, vmax=0.5)
        # mlab.outline(surf_cat1, color=(.7, .7, .7), extent=cat1_extent)
        # mlab.axes(surf_cat1, color=(.7, .7, .7), extent=cat1_extent,
        # ranges=(0, 1, 0, 1, 0, 1), xlabel='', ylabel='',
        # zlabel='Probability',
        # x_axis_visibility=False, z_axis_visibility=False)
        # mlab.text(-18, -4, '1 photon', z=-4, width=0.13)
        cat2_extent = (-4, 4, -4, 4, 0, 5)
        surf_cat2 = mlab.surf(x, y, cat2, colormap='Spectral', warp_scale=5, extent=cat2_extent, vmin=-0.5, vmax=0.5)
        # self.scene.mlab.outline(surf_cat2, color=(0.7, .7, .7), extent=cat2_extent)
        # self.scene.mlab.text(-4, -3, '2 photons', z=-4, width=0.14)
        cat3_extent = (6, 14, -4, 4, 0, 5)
        surf_cat3 = mlab.surf(x + 10, y, cat3, colormap='Spectral', warp_scale=5, extent=cat3_extent, vmin=-0.5, vmax=0.5)
        # self.scene.mlab.outline(surf_cat3, color=(.7, .7, .7), extent=cat3_extent)
        # self.scene.mlab.text(6, -2.5, '3 photons', z=-4, width=0.14)
        # self.scene.mlab.title('Multi-photons cats Wigner function')

    threevisualization1 = Instance(MlabSceneModel, args={'engine': engine})

    @observe('threevisualization1.activated')
    def update_threevisualization1(self, event=None):
        mlab = self.threevisualization1.mlab
        mlab.test_points3d()

    threevisualization2 = Instance(MlabSceneModel, args={'engine': engine})

    @observe('threevisualization2.activated')
    def update_threevisualization2(self, event=None):
        mlab = self.threevisualization2.mlab
        np.random.seed(12345)
        x = 4 * (np.random.random(500) - 0.5)
        y = 4 * (np.random.random(500) - 0.5)

        def f(x, y):
            return np.exp(-(x ** 2 + y ** 2))

        z = f(x, y)
        pts = mlab.points3d(x, y, z, z, scale_mode='none', scale_factor=0.2)
        mesh = mlab.pipeline.delaunay2d(pts)
        surf = mlab.pipeline.surface(mesh)

    testline = Instance(MlabSceneModel, args={'engine': engine})

    @observe('testline.activated')
    def update_testline(self, event=None):
        mlab = self.testline.mlab
        x=np.arange(10)
        y=np.arange(10)
        z=np.arange(10)
        mlab.plot3d(x,y,z)

    view = View(
        Tabbed(
            UItem('polydata3', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('polydata2', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('polydata1', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('boy', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('atomic_orbital', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('polydata', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('image_data', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('rectilinear_grid', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('structured_grid', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('unstructured_grid', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('delaunay_graph', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('custom_colormap', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('fluidArrow_grid', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('image_cursor_filter', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('julia_set', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('lorenz', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('RCS_3D', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('select_red_balls', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('spherical_harmonics', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('three3d_wigner', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('threevisualization1', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('threevisualization2', editor=SceneEditor(scene_class=MayaviScene)),
            UItem('testline', editor=SceneEditor(scene_class=MayaviScene)),
        ),
        height=800,
        width=1000,
        resizable=True  # We need this to resize with the parent widget
    )


from PyQt5.QtWidgets import *


class FormWidget(QWidget):
    def __init__(self):
        QWidget.__init__(self)
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        self.form = Form()
        layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)


formWidget = FormWidget()
if __name__ == '__main__':
    import sys

    app = QApplication([])
    formWidget.show()
    sys.exit(app.exec_())


